from typing import Union, Dict

import numpy as np

from numpy.fft import fft, ifft
from sklearn.utils import check_random_state
from scipy.stats import uniform


class FTSurrogate:
    def __init__(self, phase_noise_magnitude, channel_indep=False, keys=None, seed=None, prob=1.0) -> None:
        self.phase_noise_magnitude = phase_noise_magnitude
        self.channel_indep = channel_indep
        self.seed = seed
        self.prob = prob
        self.keys = keys

        self._new_random_fft_phase = {
            0: self._new_random_fft_phase_even,
            1: self._new_random_fft_phase_odd
        }

    def _new_random_fft_phase_odd(self, channels: int, length: int, seed: Union[int, None] = None) -> np.ndarray:
        rng = check_random_state(seed)
        random_phase = 2j * np.pi * rng.random((channels, (length - 1) // 2))
        return np.concatenate([np.zeros((channels, 1)), random_phase, -np.flip(random_phase, axis=-1)], axis=-1)

    def _new_random_fft_phase_even(self, channels: int, length: int, seed: Union[int, None] = None) -> np.ndarray:
        rng = check_random_state(seed)
        random_phase = 2j * np.pi * rng.random((channels, length // 2 - 1))
        return np.concatenate([np.zeros((channels, 1)), random_phase, np.zeros((channels, 1)), -np.flip(random_phase, axis=-1)], axis=-1)

    def _apply(self, sample: np.ndarray) -> np.ndarray:
        """
        Apply the Fourier transform based transformation to the sample.
        
        :param sample: Input sample to transform.
        :return: Transformed sample.
        """
        if uniform.rvs() < self.prob:
            f_transform = fft(sample, axis=-1)
            num_samples = f_transform.shape[-1]
            channels = f_transform.shape[-2] if self.channel_indep else 1
            random_phase = self._new_random_fft_phase[num_samples % 2](channels, num_samples, self.seed)

            if not self.channel_indep:
                random_phase = np.tile(random_phase, (f_transform.shape[-2], 1))

            f_shifted = f_transform * np.exp(self.phase_noise_magnitude * random_phase)
            shifted_sample = ifft(f_shifted, axis=-1)
            return shifted_sample.real.astype(np.float32)
        return sample.real.astype(np.float32)

    def __call__(self, data: Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]) -> Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]:
        if isinstance(data, dict):
            for key in self.keys:
                sample = data[key]
                data[key] = self._apply(sample)
        elif isinstance(data, np.ndarray):
            data = self._apply(data)
        return data
